"""矩阵微积分与函数模块
包含矩阵函数、矩阵导数等功能
"""
import numpy as np
from scipy.linalg import expm, logm, sqrtm, funm
from typing import Callable, Optional

class MatrixCalculus:
    @staticmethod
    def matrix_exponential(matrix: np.ndarray) -> np.ndarray:
        """
        矩阵指数函数 exp(A)
        :param matrix: 输入方阵
        :return: exp(A)
        """
        return expm(matrix)
    
    @staticmethod
    def matrix_logarithm(matrix: np.ndarray) -> np.ndarray:
        """
        矩阵对数函数 log(A)
        :param matrix: 输入方阵
        :return: log(A)
        """
        return logm(matrix)
    
    @staticmethod
    def matrix_square_root(matrix: np.ndarray) -> np.ndarray:
        """
        矩阵平方根 sqrt(A)
        :param matrix: 输入方阵
        :return: sqrt(A)
        """
        return sqrtm(matrix)
    
    @staticmethod
    def matrix_function(matrix: np.ndarray, func: Callable) -> np.ndarray:
        """
        通用矩阵函数 f(A)
        :param matrix: 输入方阵
        :param func: 标量函数
        :return: f(A)
        """
        return funm(matrix, func)
    
    @staticmethod
    def matrix_power_real(matrix: np.ndarray, p: float) -> np.ndarray:
        """
        矩阵实数幂 A^p
        :param matrix: 输入方阵
        :param p: 实数幂
        :return: A^p
        """
        eigenvals, eigenvecs = np.linalg.eig(matrix)
        powered_eigenvals = np.power(eigenvals, p)
        return eigenvecs @ np.diag(powered_eigenvals) @ np.linalg.inv(eigenvecs)
    
    @staticmethod
    def frechet_derivative(matrix: np.ndarray, direction: np.ndarray, 
                          func: Callable, h: float = 1e-8) -> np.ndarray:
        """
        Fréchet导数
        :param matrix: 基点矩阵
        :param direction: 方向矩阵
        :param func: 矩阵函数
        :param h: 步长
        :return: Fréchet导数
        """
        return (func(matrix + h * direction) - func(matrix)) / h
    
    @staticmethod
    def matrix_taylor_expansion(matrix: np.ndarray, func: Callable, 
                              order: int = 5) -> np.ndarray:
        """
        矩阵函数的Taylor展开
        :param matrix: 输入矩阵
        :param func: 标量函数
        :param order: 展开阶数
        :return: Taylor展开近似
        """
        n = matrix.shape[0]
        result = np.eye(n) * func(0)  # 常数项
        
        # 计算各阶导数（数值近似）
        h = 1e-6
        matrix_power = np.eye(n)
        
        for k in range(1, order + 1):
            matrix_power = matrix_power @ matrix
            
            # 数值计算k阶导数
            if k == 1:
                derivative = (func(h) - func(0)) / h
            else:
                # 使用有限差分近似高阶导数
                derivative = 1.0  # 简化实现
            
            result += derivative * matrix_power / np.math.factorial(k)
        
        return result
    
    @staticmethod
    def jacobian_matrix(func: Callable, point: np.ndarray, h: float = 1e-8) -> np.ndarray:
        """
        计算向量函数的雅可比矩阵
        :param func: 向量函数
        :param point: 计算点
        :param h: 步长
        :return: 雅可比矩阵
        """
        n = len(point)
        f0 = func(point)
        m = len(f0)
        
        jacobian = np.zeros((m, n))
        
        for j in range(n):
            point_plus = point.copy()
            point_plus[j] += h
            f_plus = func(point_plus)
            
            jacobian[:, j] = (f_plus - f0) / h
        
        return jacobian
    
    @staticmethod
    def hessian_matrix(func: Callable, point: np.ndarray, h: float = 1e-6) -> np.ndarray:
        """
        计算标量函数的Hessian矩阵
        :param func: 标量函数
        :param point: 计算点
        :param h: 步长
        :return: Hessian矩阵
        """
        n = len(point)
        hessian = np.zeros((n, n))
        
        for i in range(n):
            for j in range(n):
                # 计算二阶偏导数
                point_ij = point.copy()
                point_ij[i] += h
                point_ij[j] += h
                
                point_i = point.copy()
                point_i[i] += h
                
                point_j = point.copy()
                point_j[j] += h
                
                hessian[i, j] = (func(point_ij) - func(point_i) - func(point_j) + func(point)) / (h * h)
        
        return hessian